package com.simple.llm.service.impl;

import com.github.zuihou.base.R;
import com.simple.llm.config.LLMProperties;
import com.simple.llm.domain.dto.LLMStreamResultDTO;
import com.simple.llm.domain.dto.NormalChatDTO;
import com.simple.llm.domain.dto.StreamChatDTO;
import com.simple.llm.domain.vo.LLMResponseVO;
import com.simple.llm.service.LLMStrategy;
import com.volcengine.ark.runtime.model.Usage;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.servlet.http.HttpServletResponse;
import java.util.List;
import java.util.Locale;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
public class DouBaoLLMStrategyImpl implements LLMStrategy {
    private final LLMProperties.GptProperties gptProperties;
    private final ArkService service;

    public DouBaoLLMStrategyImpl(LLMProperties.GptProperties gptProperties) {
        this.gptProperties = gptProperties;
        service = ArkService.builder().apiKey(gptProperties.getApiKeys().get(0)).build();
    }

    @Override
    public R<LLMResponseVO> chat(NormalChatDTO normalChatDTO) {
        List<ChatMessage> messages = normalChatDTO.getMessages().stream().map(o ->ChatMessage.builder()
                .content(o.getContent())
                .role(ChatMessageRole.valueOf(o.getRole().toUpperCase(Locale.ENGLISH)))
                .build()).collect(Collectors.toList());

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(gptProperties.getSpecialKey())
                .messages(messages)
                .stream(Boolean.FALSE)
                .maxTokens(gptProperties.getMaxContextTokens())
                .temperature((double)(normalChatDTO.getTemperature() == null ? 0.8F : normalChatDTO.getTemperature()))
                .topP((double)(normalChatDTO.getTopP() == null ? 0.75F : normalChatDTO.getTopP()))
                .build();

        try {
            ChatCompletionResult chatCompletion = service.createChatCompletion(chatCompletionRequest);
            if(chatCompletion != null){
                Usage usage = chatCompletion.getUsage();
                String content = chatCompletion.getChoices().get(0).getMessage().getContent().toString();
                LLMResponseVO vo = LLMResponseVO.builder()
                        .content(content)
                        .inputTokenCount(usage.getPromptTokens())
                        .outputTokenCount(usage.getCompletionTokens())
                        .desc(normalChatDTO.getDesc())
                        .createUser(normalChatDTO.getCreateUser())
                        .build();
                return R.success(vo);
            }
        }catch (Exception exception){
            log.error("调用豆包异常",exception);
        }
        return R.fail("调用异常");
    }

    @Override
    public ResponseBodyEmitter streamChat(StreamChatDTO streamChatDTO, HttpServletResponse response, Consumer<LLMStreamResultDTO> resultHandler) {
        return null;
    }
}
